%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% GE_IRF_Baseline.m
% Simulate a GE heterogeneous firms investment model with macro shocks and real and financial adjustment costs 
%
% Ivan Alfaro, Nick Bloom and Xiaoji Lin 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


clc; clearvars;

load ks_het_firms_baseline.mat; 

%control the IRF
numper          = 2000;  % period for initial point of K using the simulation of the model simulation 
numsimIRF       = 500;   % number of shocked economies to simulate
lengthIRF       = 30;    % length of each economy
shockperIRF     = 200;   % period in which to shock the economy
checkbounds     = 1;

numperIRF       = numsimIRF*lengthIRF; %total number of periods in IRF simulation
Ttot            = lengthIRF + shockperIRF;
maxpit          = 50;
pwindow         = 0.1;
pcutoff         = 15;
sigmaAdummy     = 1;


rng('default');                                  % fix the seed 
rng(1);

%then, initialize the aggregate series to 0
YsimIRF    = zeros(Ttot,numsimIRF); KZsimIRF = YsimIRF;    
KsimIRF    = YsimIRF;               NsimIRF = YsimIRF; IsimIRF    = YsimIRF;       
ACsimIRF   = YsimIRF;               CsimIRF = YsimIRF; psimIRF    = YsimIRF;       
DivsimIRF  = YsimIRF;               HsimIRF = YsimIRF; perrsimIRF = YsimIRF;
ECsimIRF   = YsimIRF;   	      EsimIRF   = YsimIRF; 		

%start with aggregate capital guessed at some reasonable value from uncond simulation
KsimIRF(1,:) = Ksim(numper);                %Kinit K0(ceil(Knum/2));%
pmat       = permute(repmat(p0',[1 knum,nnum,znum,snum,knum,nnum]),[2 3 4 5 6 7 1]);%(k',n',z,s,k,n,p)
z1         = permute(repmat(z0,[1 snum knum nnum knum nnum   pnum]),[5 6 1 2 3 4 7]); %(k',n',z,s,k,n,p)
k1         = permute(repmat(k0,[1 znum snum nnum knum nnum   pnum]),[5 6 2 3 1 4 7]); %(k',n',z,s,k,n,p)
kprime1    = permute(repmat(k0,[1 znum snum knum nnum nnum   pnum]),[1 6 2 3 4 5 7]); %(k',n',z,s,k,n,p)
n1         = permute(repmat(n0',[1 znum snum knum knum nnum  pnum]),[5 6 2 3 4 1 7]); %(k',n',z,s,k,n,p)
nprime1    = permute(repmat(n0',[1 znum snum knum nnum knum  pnum]),[6 1 2 3 4 5 7]); %(k',n',z,s,k,n,p)
k1_z_alpha = z1.*k1.^alpha;                 %precomputed output
ivaltemp = kprime1 - (1-delta)*k1;                            %(k',n',z,s,k,n,p)
ivaltemp(abs(ivaltemp)<=ZeroBoundary) =0;
ival     = ivaltemp;
hval     = nprime1 - (1+rf)*kappa*n1;                         %(k',n',z,s,k,n,p)  

%%%%%%%%%
%%%%simulation
%%%%%%%%%
%exogenous macro shocks first

AsimposIRF         = zeros(Ttot,numsimIRF);
AsimposIRF(1,:)    = Ainit;
AshocksIRF         = rand(Ttot,numsimIRF);  %repmat(Ashocks,[1 numsimIRF]);
sigmaAsimposIRF    = zeros(Ttot,numsimIRF);
sigmaAsimposIRF(1,:) = Sinit;

sigmaAshocksIRF    = rand(Ttot,numsimIRF);  %repmat(sigmaAshocks,[1 numsimIRF]); 
FsimposIRF         = ones(Ttot,numsimIRF);
FsimposIRF(1,:)    = Finit;
FshocksIRF         = rand(Ttot,numsimIRF); 

%%
for simct = 1:numsimIRF

    %first, initialize the distribution over endogenous variables
    distzsknIRF             = zeros(znum,snum,knum,nnum,Ttot);
    distzsknIRF(:,:,:,:,1)  = distzskn(:,:,:,:,numper); 
    distzsknIRF(:,:,:,:,1)  = distzsknIRF(:,:,:,:,1)/ sum(sum(sum(sum(distzsknIRF(:,:,:,:,1)))));

     for t =1:Ttot-1
            if t==shockperIRF
                if sigmaAdummy ==1
                         sigmaAsimposIRF(t,simct) = 2;
                         Sct      = sigmaAsimposIRF(t,simct); 
                         sprimect = find(PisumSigma(Sct,:) >= sigmaAshocksIRF(t+1,simct), 1); %     %compare today's uniform shock to transition matrix intervals
                         sigmaAsimposIRF(t+1,simct) = sprimect;        %aggregate volatility %     %store simulated value
                         % agg prodcutivity
                         Act       = AsimposIRF(t,simct);
                         Aprimect  = find(PisumA(Act,:,Sct) >= AshocksIRF(t+1,simct), 1);  %     %compare today's uniform shock to transition matrix intervals
                         AsimposIRF(t+1,simct) = Aprimect;              %store simulated value
                         % financial shock
                         Fct       = FsimposIRF(t,simct);
                         Fprimect  = find(PisumF(Fct,:) >= FshocksIRF(t+1,simct), 1); %            %compare today's uniform shock to transition matrix intervals
                         FsimposIRF(t+1,simct)= Fprimect;        
                end
            else
                         % time-varying uncertainty evolve normally
                         Sct       = sigmaAsimposIRF(t,simct);
                         sprimect  = find(PisumSigma(Sct,:) >= sigmaAshocksIRF(t+1,simct), 1); %     %compare today's uniform shock to transition matrix intervals
                         sigmaAsimposIRF(t+1,simct)= sprimect;        %aggregate volatility %     %store simulated value
                         % agg prodcutivity
                         Act       = AsimposIRF(t,simct);
                         Aprimect  = find(PisumA(Act,:,Sct) >= AshocksIRF(t+1,simct), 1);  %     %compare today's uniform shock to transition matrix intervals
                         AsimposIRF(t+1,simct) = Aprimect;    %     %store simulated value
                         % financial shock
                         Fct       = FsimposIRF(t,simct);
                         Fprimect  = find(PisumF(Fct,:) >= FshocksIRF(t+1,simct), 1); %            %compare today's uniform shock to transition matrix intervals
                         FsimposIRF(t+1,simct)= Fprimect;                                          %store simulated value  

            end
     end
    
% disp('Simulating the model for impulse response functions.')

    Act = 0;Sct = 0;Fct =0;
    for t=1:(Ttot-1)
         %extract macro prod
            Act       = AsimposIRF(t,simct);
            Aval      = A0(Act);
            Sct       = sigmaAsimposIRF(t,simct);
            sigmaAval = sigmaA(Sct);
            Fct       = FsimposIRF(t,simct);
            FSval     = FinC(Fct);     % financial cost shock
        %extract macro capital and forecasts
            Kval      = KsimIRF(t,simct);
            Kval      = min(max(Kval,Kmin),Kmax);
            Kprimeval = exp(thetaKold(Act,Sct,Fct,1)+ thetaKold(Act,Sct,Fct,2)*log(Kval));
            Kprimeind = sum(Kprimeval>=K0);          
            
            % guard against off grid point values
            if (Kprimeind==Knum) 
                Kprimeind = Knum-1;
                Kprimewgt = 1.0;
            elseif (Kprimeind==0)
                Kprimeind = 1;
                Kprimewgt = 0.0;
            else
               Kprimewgt  = (Kprimeval-K0(Kprimeind))/(K0(Kprimeind+1)-K0(Kprimeind));
            end 
            pfcstval    = exp(thetapold(Act,Sct,Fct,1) + thetapold(Act,Sct,Fct,2)*log(Kval));

          kz       = z1.*k1;
          yval     = Aval*k1_z_alpha;                                   %(k',n',z,s,k,n,p)
          ACval    = c_k*(ival~=0).*yval;                               %(k',n',z,s,k,n,p)
          Eval     = yval - ival - ACval -hval;                         %(k',n',z,s,k,n,p)
          ECval    = FSval.*abs(Eval).*(Eval<0);                        %(k',n',z,s,k,n,p)
          Rmatp    = pmat.*(Eval-ECval);                                %(k',n',z,s,k,n,p)
        
        if (t==shockperIRF) 
                if (sigmaAdummy ==1)
                         sct_l = 1; sct_h = 2;
                         distzsknIRF_new = zeros(znum,snum,knum,nnum);
                         for zct=1:znum %z
                             for kct=1:knum %k     
                                 for nct = 1:nnum %n
                                     if distzsknIRF(zct,sct_l,kct,nct,t)>0
                                         %move the sigma_l distribution to sigma_h; the new
                                         %distribution has a mass of 1 in sigma_h, a mass of 0 in sigma_l
                                         distzsknIRF_new(zct,sct_h,kct,nct) = distzsknIRF(zct,sct_l,kct,nct,t)+...
                                             distzsknIRF(zct,sct_h,kct,nct,t); 
                                     end
                                 end
                             end
                         end

                         disttemp = repmat(distzsknIRF_new,[1 1 1 1 pnum]);
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
                        % compute the conditional expectation
                        EVMAT    = zeros(knum,nnum,znum);
                        temp = (1-Kprimewgt)*squeeze(V(:,:,:,Kprimeind,:,:,:,:))...
                                + Kprimewgt*squeeze(V(:,:,:,Kprimeind+1,:,:,:,:));             %(A',S',F',z',s',k',n')                                
                        temp1 = PiF(Fct,:)*reshape(PiSigma(Sct,:)*reshape(PiA(Act,:,Sct)*reshape(temp,...
                                        [Anum,Snum*Fnum*znum*snum*knum*nnum]),...
                                        [Snum,Fnum*znum*snum*knum*nnum]),...
                                        [Fnum,znum*snum*knum*nnum]);  
                        for zct = 1:znum 
                                EVMAT(:,:,zct) = reshape(Pisigma(sct_h,:)*reshape(Piz(zct,:,sct_h)*reshape(temp1,[znum,snum*knum*nnum]),[snum,knum*nnum]),[knum nnum]); 
                                            
                        end
            
                        EVMAT_H = repmat(EVMAT,[1 1 1 snum knum nnum pnum]);                   %(k',n',z,s,k,n,p)
           
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
                        yval2   = repmat(yval(:,:,:,sct_h,:,:,:),[1 1 1 snum 1 1 1]);            %(k',n',z,s,k,n,p)
                        RHSMAT2 = Rmatp + beta*EVMAT_H;
                        RHSMAT2 = repmat(RHSMAT2(:,:,:,sct_h,:,:,:),[1 1 1 snum 1 1 1]);        %(k',n',z,s,k,n,p)
                        
                        [~,kprimeinddum] = max(max(RHSMAT2,[],2),[],1);
                        [~,nprimeinddum] = max(max(RHSMAT2,[],1),[],2);
                        kprimeindp = squeeze(kprimeinddum);
                        nprimeindp = squeeze(nprimeinddum);
                        kprimep    = k0(kprimeindp);
                        nprimep    = n0(nprimeindp);
                        
                        yvalp = squeeze(yval2(1,1,:,:,:,:,:));                           %(z,s,k,n,p)
                        kzp   = squeeze(kz(1,1,:,:,:,:,:));                                %(z,s,k,n,p)
                        ivalptemp = (kprimep - (1-delta)*squeeze(k1(1,1,:,:,:,:,:)));
                        ivalptemp(abs(ivalptemp)<=ZeroBoundary)=0; 
                        ivalp  = repmat(ivalptemp(:,sct_h,:,:,:),[1 snum 1 1 1]);        %(z,s,k,n,p)
                        acvalp = c_k*(ivalp~=0).*yvalp;
                        hvalp  = nprimep-(1+rf)*kappa*squeeze(n1(1,1,:,:,:,:,:));
                        hvalp  = repmat(hvalp(:,sct_h,:,:,:),[1 snum 1 1 1]);            %(z,s,k,n,p)
                        evalp  = (yvalp-ivalp-acvalp-hvalp);
                        ecvalp = FSval*abs(evalp).*(evalp<0);  %(z,s,k,n,p)
                        divvalp  = evalp-ecvalp;                                         %(z,s,k,n,p)
                        
                        kprimep  = repmat(kprimep(:,sct_h,:,:,:),[1 snum 1 1 1]);        %(z,s,k,n,p)
                        nprimep  = repmat(nprimep(:,sct_h,:,:,:),[1 snum 1 1 1]);        %(z,s,k,n,p)

                       Yp0        = squeeze(sum(sum(sum(sum(yvalp.*disttemp)))));
                       Ip0        = squeeze(sum(sum(sum(sum(ivalp.*disttemp)))));
                       KZp0       = squeeze(sum(sum(sum(sum(kzp.*disttemp)))));
                       Hp0        = squeeze(sum(sum(sum(sum(hvalp.*disttemp)))));
                       Kbarprimep0= squeeze(sum(sum(sum(sum(kprimep.*disttemp)))));
                       Nbarprimep0= squeeze(sum(sum(sum(sum(nprimep.*disttemp)))));
                       Divp0      = squeeze(sum(sum(sum(sum(divvalp.*disttemp)))));
                       ACp0       = squeeze(sum(sum(sum(sum(acvalp.*disttemp)))));
                       ECp0       = squeeze(sum(sum(sum(sum(ecvalp.*disttemp)))));

                       polmatp_interp_k = repmat(kprimeindp(:,sct_h,:,:,:),[1 snum 1 1 1]);  %(z,s,k,n,p)
                       polmatp_interp_n = repmat(nprimeindp(:,sct_h,:,:,:),[1 snum 1 1 1]); %(z,s,k,n,p)
              
                end  %(sigmaAdummy ==1 )
            
        else % normal conditional expectation
                EVMAT  = zeros(knum,nnum,znum,sct);
                temp   = (1-Kprimewgt)*squeeze(V(:,:,:,Kprimeind,:,:,:,:))...
                          + Kprimewgt*squeeze(V(:,:,:,Kprimeind+1,:,:,:,:));             %(A',S',F',z',s',k',n')                                
                temp1 = PiF(Fct,:)*reshape(PiSigma(Sct,:)*reshape(PiA(Act,:,Sct)*reshape(temp,...
                                [Anum,Snum*Fnum*znum*snum*knum*nnum]),...
                                [Snum,Fnum*znum*snum*knum*nnum]),...
                                [Fnum,znum*snum*knum*nnum]);  %(1,S',F',z',s',k',n') to %(S',F'*z'*s'*k'*n') to (F',z'*s'*k'*n') to (z',s',k',n')
                for zct = 1:znum 
                    for sct =1:snum   
                        EVMAT(:,:,zct,sct) = reshape(Pisigma(sct,:)*reshape(Piz(zct,:,sct)*reshape(temp1,[znum,snum*knum*nnum]),[snum,knum*nnum]),[knum nnum]); 
                                    %(z',s',k') to (z',s'*k') to (1,s',k') to    %(k',1)
                     end
                end
            
              EVMAT_H = repmat(EVMAT,[1 1 1 1 knum nnum pnum]);               %(k',n',z,s,k,n,p)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
              disttemp= repmat(distzsknIRF(:,:,:,:,t),[1 1 1 1 pnum]);%(z,s,k,n,p)
              RHSMAT2 = Rmatp + beta*EVMAT_H;
              [~,kprimeinddum] = max(max(RHSMAT2,[],2),[],1);
              [~,nprimeinddum] = max(max(RHSMAT2,[],1),[],2);
              kprimeindp = squeeze(kprimeinddum);
              nprimeindp = squeeze(nprimeinddum);
              kprimep    = k0(kprimeindp);
              nprimep    = n0(nprimeindp);

               yvalp = squeeze(yval(1,1,:,:,:,:,:));
               kzp   = squeeze(kz(1,1,:,:,:,:,:));                           %(z,s,k,n,p)
               ivalptemp = (kprimep - (1-delta)*squeeze(k1(1,1,:,:,:,:,:)));
               ivalptemp(abs(ivalptemp)<=ZeroBoundary)=0; 
               ivalp = ivalptemp;
               acvalp = c_k*(ivalp~=0).*yvalp;
               hvalp  = nprimep-(1+rf)*kappa*squeeze(n1(1,1,:,:,:,:,:));
               evalp  = (yvalp-ivalp-acvalp-hvalp);
               ecvalp = FSval*abs(evalp).*(evalp<0);  %(z,s,k,n,p)
               divvalp = evalp-ecvalp;

               Yp0        = squeeze(sum(sum(sum(sum(yvalp.*disttemp)))));   
               KZp0       = squeeze(sum(sum(sum(sum(kzp.*disttemp)))));
               Ip0        = squeeze(sum(sum(sum(sum(ivalp.*disttemp)))));
               Hp0        = squeeze(sum(sum(sum(sum(hvalp.*disttemp)))));
               Kbarprimep0= squeeze(sum(sum(sum(sum(kprimep.*disttemp)))));
               Nbarprimep0= squeeze(sum(sum(sum(sum(nprimep.*disttemp)))));
               Divp0      = squeeze(sum(sum(sum(sum(divvalp.*disttemp)))));
               ACp0       = squeeze(sum(sum(sum(sum(acvalp.*disttemp)))));
               ECp0       = squeeze(sum(sum(sum(sum(ecvalp.*disttemp)))));
               polmatp_interp_k = kprimeindp; %(z,s,k,n,p)
               polmatp_interp_n = nprimeindp; %(z,s,k,n,p)
         
        end
        
           Cvalp = Yp0-Ip0-ACp0-ECp0;
           Cp0   = Cvalp;
           ep0   = 1./p0'-Cvalp;           
           
          %set up the boundaries of the bisection
            pvala = plb;
            pvalb = pub;
            pvalc = pvala + (0.67)*(pvalb-pvala);

            %iterate over the market-clearing value of p, for each value
            %redoing the optimization of the RHS of the Bellman equation
            %and recomputing policies
        for piter=1:maxpit
            %initialize the price-dependent policies to 0
    %         kprimeinddum(:,:) = 0; kprimeindp(:,:) = 0; 

            %brent setup
            if (piter==1); pval = pvala;end
            if (piter==2); pval = pvalb;end
            if (piter==3); pval = pvalc;end

            %brent restart
            if (piter==pcutoff)
                pvala = pfcstsim(t)-4.0*pwindow;
                pval = pvala;
            elseif (piter==pcutoff+1) 
                pvalb =  pfcstsim(t)+4.0*pwindow;
                pval = pvalb;
            elseif (piter==pcutoff+2)
                pvalc = pvala + (0.67) * (pvalb-pvala); 
                pval = pvalc;
            end 

            %if not restarting or initializing
            if ((piter>3 && piter<pcutoff) || (piter>pcutoff+2)) 
                %first, try inverse quadratic interpolation of the excess demand function
                pval = ( pvala * fb * fc ) / ( (fa - fb) * (fa - fc) ) ...
                    + ( pvalb * fa * fc ) / ( (fb - fa) * (fb - fc ) ) ...
                    + ( pvalc * fa * fb ) / ( (fc - fa) * (fc - fb ) );
                %if it lies within bounds, and isn't too close to the bounds, then done
                %o/w, take bisection step
                if (min(abs(pvala - pval),abs(pvalb-pval))<abs((pvalb-pvala)/9)) || (pval<pvala) || (pval>pvalb)
    %                 ((minval( (/ abs(pvala - pval), abs(pvalb-pval) /) )<&
    %                     abs( (pvalb-pvala)/dble(9.0) ) ).or.(pval<pvala).or.(pval>pvalb))   then
                    pval = (pvala + pvalb) /2;
                end
            end

            %actually evaluate consumption approximation function
            pval = min(max(p0(1),pval),p0(pnum));         %minval( (/ maxval( (/ p0(1) , pval /) ) , p0(pnum) /) )   
            pind = sum(pval>=p0);%pind = pnum/2;  call hunt(p0,pnum,pval,pind)
    % 	    pwgt = (pval - p0(pind)) / (p0(pind + 1) - p0(pind));

            % guard against off grid point values
            if (pind==pnum) 
                pind = pnum-1;
                pwgt = 1.0;
            elseif (pind==0)
                pind = 1;
                pwgt = 0.0;
            else
                pwgt = (pval - p0(pind)) / (p0(pind + 1) - p0(pind));
            end 

            %actually perform linear interpolation of the consumption function
            CvalpNew = Cp0(pind)*(1-pwgt) + Cp0(pind+1)*pwgt;
            CvalpNew = max(CvalpNew,kmin);
            
            %are you initializing the brent?
            if (piter==1); fa = (1/pval) - CvalpNew;end
            if (piter==2); fb = (1/pval) - CvalpNew;end
            if (piter==3); fc = (1/pval) - CvalpNew;end

            %are you restarting the brent?
            if (piter==pcutoff);   fa = (1/pval) - CvalpNew; end
            if (piter==pcutoff+1); fb = (1/pval) - CvalpNew; end
            if (piter==pcutoff+2); fc = (1/pval) - CvalpNew; end

            %what is the error given by this implied consumption?
            perror = (1/pval) - CvalpNew;

            %if not restarting or initializing
            if ((piter>3 && piter<pcutoff)||(piter>pcutoff+2))  
                if (perror<0) 
                    pvalc = pvalb; fc = fb;
                    pvalb = pval; fb = perror;
                    %pval a doesn't change
                elseif (perror>=0) 
                    pvalc = pvala; fc = fa;
                    pvala = pval; fa = perror;
                    %pval b doesn't change
                end
            end

            %exit criterion for market-clearing
            perror = log(pval*CvalpNew);
            if (abs(perror)<perrortol && piter>2) 
                break
            end  %piter
        end
              %insert market-clearing price and other linearly interpolated stuff into sim series
            perrsimIRF(t,simct) = perror;
            psimIRF(t,simct)    = pval; %this is the most recently run p from the clearing algorithm
            CsimIRF(t,simct)    = CvalpNew; %this is already the linearly interpolated consumption C
            YsimIRF(t,simct)    = Yp0(pind)*(1.0-pwgt) + Yp0(pind+1)*pwgt; %linearly interpolated output Y
            IsimIRF(t,simct)    = Ip0(pind)*(1.0-pwgt) + Ip0(pind+1)*pwgt; %linearly interpolated investment I
            HsimIRF(t,simct)    = Hp0(pind)*(1.0-pwgt) + Hp0(pind+1)*pwgt; %linearly interpolated investment H
            ACsimIRF(t,simct)   = ACp0(pind)*(1.0-pwgt)+ ACp0(pind+1)*pwgt; %linearly interpolated ACk 
            KsimIRF(t+1,simct)  = Kbarprimep0(pind)*(1.0-pwgt) + Kbarprimep0(pind+1)*pwgt; %linearly interpolated K' 
            NsimIRF(t+1,simct)  = Nbarprimep0(pind)*(1.0-pwgt) + Nbarprimep0(pind+1)*pwgt; %linearly interpolated cash N'
            DivsimIRF(t,simct)  = Divp0(pind)*(1.0-pwgt) + Divp0(pind+1)*pwgt; %linearly interpolated investment Div
            ECsimIRF(t,simct)   = ECp0(pind)*(1.0-pwgt) + ECp0(pind+1)*pwgt;   %linearly interpolated EC
            KZsimIRF(t,simct)   = KZp0(pind)*(1.0-pwgt) + KZp0(pind+1)*pwgt;  %linearly interpolated KZ 
            
            %output price stats on certain periods
          if (mod(t,50)==1)
                disp(['IRF = ',num2str(simct), 't = ',num2str(t), ' iter ',num2str(piter),' p ',num2str(pval),' err ',num2str(perror), ' A ',num2str(A0(Act)),' Kprime ' ,num2str(Ksim(t+1)),...
                    ' Nprime ' ,num2str(Nsim(t+1))])        
          end
      
          %if K hits boundary, reset the next period distribution
          if  KsimIRF(t+1,simct)<=kmin ||  KsimIRF(t+1,simct)>=kmax      
                    KsimIRF(t+1,simct)   = Ksim(numper);
                    distzsknIRF(:,:,:,:,t+1)= distzskn(:,:,:,:,numper);%distoldIRF; 
         end
          
           %update the t+1 distribution
          if t==shockperIRF   
                  if (sigmaAdummy ==1 )
                         sct_h = 2;
                         for zct=1:znum
                              for sct =1:snum
                                    for kct=1:knum 
                                        for nct=1:nnum 
                                            if (distzsknIRF_new(zct,sct,kct,nct)>disttol)       
                                                %based on the latest price, what is the policy here at pind?
                                                polstar_k_0 = polmatp_interp_k(zct,sct,kct,nct,pind);
                                                polstar_n_0 = polmatp_interp_n(zct,sct,kct,nct,pind);

                                                %insert distributional weight in appropriate slots next period
                                                distzsknIRF(:,:,polstar_k_0,polstar_n_0,t+1) = distzsknIRF(:,:,polstar_k_0,polstar_n_0,t+1) + ...
                                                    repmat(Piz(zct,:,sct_h)',1,snum).*repmat(Pisigma(sct_h,:),znum,1)*distzsknIRF_new(zct,sct,kct,nct) * (1.0-pwgt);

                                                %based on the latest price, what is the policy here at pind+1?
                                                polstar_k_1 = polmatp_interp_k(zct,sct,kct,nct,pind+1);
                                                polstar_n_1 = polmatp_interp_n(zct,sct,kct,nct,pind+1);

                                                %insert distributional weight in appropriate slots next period
                                                distzsknIRF(:,:,polstar_k_1,polstar_n_1,t+1) = distzsknIRF(:,:,polstar_k_1,polstar_n_1,t+1) + ...
                                                    repmat(Piz(zct,:,sct_h)',1,snum).*repmat(Pisigma(sct_h,:),znum,1)*distzsknIRF_new(zct,sct,kct,nct) * pwgt;
                                            end
                                        end %nct
                                    end  %kct
                              end %sct
                         end  %zct      
                    
                  end % if SigmaDummy
  
          else % t is not shockperIRF
                      for zct=1:znum
                           for sct =1:snum
                                for kct=1:knum 
                                     for nct=1:nnum 
                                           if (distzsknIRF(zct,sct,kct,nct,t)>disttol)       
                                                    %based on the latest price, what is the policy here at pind?
                                                    polstar_k_0 = polmatp_interp_k(zct,sct,kct,nct,pind);
                                                    polstar_n_0 = polmatp_interp_n(zct,sct,kct,nct,pind);

                                                    %insert distributional weight in appropriate slots next period
                                                    distzsknIRF(:,:,polstar_k_0,polstar_n_0,t+1) = distzsknIRF(:,:,polstar_k_0,polstar_n_0,t+1) + ...
                                                        repmat(Piz(zct,:,sct)',1,snum).*repmat(Pisigma(sct,:),znum,1)*distzsknIRF(zct,sct,kct,nct,t) * (1.0-pwgt);

                                                    %based on the latest price, what is the policy here at pind+1?
                                                    polstar_k_1 = polmatp_interp_k(zct,sct,kct,nct,pind+1);
                                                    polstar_n_1 = polmatp_interp_n(zct,sct,kct,nct,pind+1);

                                                    %insert distributional weight in appropriate slots next period
                                                    distzsknIRF(:,:,polstar_k_1,polstar_n_1,t+1) = distzsknIRF(:,:,polstar_k_1,polstar_n_1,t+1) + ...
                                                        repmat(Piz(zct,:,sct)',1,snum).*repmat(Pisigma(sct,:),znum,1)*distzsknIRF(zct,sct,kct,nct,t) * pwgt;
                                                
                                           end
                                     end %nct
                                end  %kct
                           end %sct
                        end  %zct       
          end %if t==shockperIRF  
          
          %now, round to make sure that you're ending up with a distribution which makes sense each period
          distzsknIRF(:,:,:,:,t+1) = distzsknIRF(:,:,:,:,t+1)./sum(sum(sum(sum(distzsknIRF(:,:,:,:,t+1)))));     
              
    end % t
   
end %simct

TFPsimIRF = KZsimIRF./KsimIRF;

clear ACval acvalp Cvalp Cvalp divval divvalp EVMAT EVMAT_H EVMAT_H3 ival ivalp ivalptemp ivaltemp k1 k1_ly
clear kprime1 kprime_ind Kprimeeest kprimeinddum kprimep nval nvalp p0mat p0z_l p0z_y p1 yval yvalp distzskn
clear H hval hvalp kprimep kprimeindp Kprimevalmat KRMSEstore lval lvalp n1 nprime nprime1 nprimeindum nprimeindp nprimen nprimep
clear polmatp_interp_k polmatp_interp_n pR2store pRMSEstore temp tempz tempzs 
clear RHSMAT2 Rmatp nprimeinddum KR2changestore KR2sotre kprime Eval evalp Eval ecvalp ECval ECval ECp0 disttemp

if sigmaAdummy == 1
    save IRF_Sigma_Baseline.mat
end


